package com.feidee.fd.sml.algorithm.component.feature

import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.feature.StringIndexer
import org.apache.spark.sql.DataFrame

/**
  * @Author tangjinyuan
  * @Date 2019/03/21 16:33
  * @Description StringIndexEncoder 字符串转索引 特征处理
  * @Reviewer dongguosheng
  */
case class StringIndexEncoderParam(
                                    override val input_pt: String,
                                    override val output_pt: String,
                                    override val hive_table: String,
                                    override val flow_time: String,
                                    override val inputCol: String,
                                    override val outputCol: String,
                                    override val preserveCols: String,
                                    override val modelPath: String,
                                    // 默认值error,对空缺数值的处理规定，有如下参数：“skip” - 过滤掉此条数据；“error” - 抛出错误；“keep” - 对其设置一个新的索引值
                                    handleInvalid: String
                                  ) extends FeatureParam {

  def this() = this(null, null, null, null, "input", "features", null, null, "error")

  override def verify(): Unit = {
    super.verify()
    require(Set("skip", "error", "keep").contains(handleInvalid.toLowerCase),"handleInvalid must be in (skip,error,keep)")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("handleInvalid" -> handleInvalid)
    map
  }
}


class StringIndexEncoder extends AbstractFeatureEncoder[StringIndexEncoderParam] {

  override def setUp(param: StringIndexEncoderParam, data: DataFrame): Array[PipelineStage] = {

    val stringIndex = new StringIndexer()
      .setHandleInvalid(param.handleInvalid)
      .setInputCol(param.inputCol)
      .setOutputCol(param.outputCol)

    Array(stringIndex)
  }

}

object StringIndexEncoder {

  def apply(paramStr: String): Unit = {
    new StringIndexEncoder()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    StringIndexEncoder(args(0))
  }
}
